import httpcore, json, strutils, strformat, parseopt, asyncdispatch, asynchttpserver, os, osproc, tables
type
    Config* = object
        port*: uint
        whitelistFile*: string


type
    RouteHandler = proc(req: Request): Future[(JsonNode, HttpCode)] {.gcsafe.}

var routes {.threadvar.}: Table[string, Table[string, RouteHandler]]
var cfg {.threadvar.}: Config

proc getConfig(): Config =
    var cfg = Config(port: 9001, whitelistFile: "whitelist.txt")
    for kind, key, val in getopt():
        case kind
        of cmdArgument:
            raise newException(ValueError, "invalid argument")
        of cmdLongOption, cmdShortOption:
            case key
            of "port", "p": cfg.port = parseUint(val)
            of "whitelist", "w": cfg.whitelistFile = val
        of cmdEnd: raise newException(ValueError, "invalid argument")
    return cfg


proc isValidIp(ip: string): bool =
    let parts = ip.split(".")
    if parts.len != 4: return false
    for part in parts:
        if part.len == 0 or part.len > 3:
            return false
        for c in part:
            if not c.isDigit:
                return false
        if part.parseInt > 255:
            return false
    return true

proc isWhitelisted(ip: string, whitelistFile: string): bool =
    if fileExists(whitelistFile):
        for line in lines(whitelistFile):
            let strippedLine = line.strip()
            if strippedLine.len > 0 and ip.startsWith(strippedLine):
                return true
    return false

proc addToIpset(ip: string, timeout: int): bool =
    let sh = &"/usr/sbin/ipset -exist add blacklist {ip} timeout {timeout}"
    return execCmd(sh) == 0

proc block_ip(req: Request): Future[(JsonNode, HttpCode)] {.async, gcsafe.} =
    var ip = ""
    var t = 0
    try:
        let jsonResult = parseJson(req.body)
        if jsonResult.hasKey("ip") and jsonResult.hasKey("t"):
            ip = jsonResult["ip"].getStr
            t = jsonResult["t"].getInt
        else:
            return (%*{"error": "Invalid JSON body"}, Http400)
    except JsonParsingError:
        return (%*{"error": "Invalid JSON body"}, Http400)
    if t < 1 or t > 2592000:
        return (%*{"error": "Invalid time"}, Http400)
    if isWhitelisted(ip, cfg.whitelistFile):
        return (%*{"error": "IP is whitelisted"}, Http200)
    if not isValidIp(ip):
        if isValidIp(&"{ip}.0"):
            ip = &"{ip}.0/24"
        else:
            return (%*{"error": "Invalid IP address"}, Http400)
    else:
        if ip.endsWith(".0"):
            ip = &"{ip}/24"
    if addToIpset(ip, t):
        return (%*{"ok": "IP blocked successfully"}, Http200)
    else:
        return (%*{"error": "Failed to block IP"}, Http500)


proc route(m: string, path: string, handler: RouteHandler) =
    if not routes.hasKey(path):
        routes[path] = initTable[string, RouteHandler]()
    routes[path][m] = handler

proc cb(req: Request) {.async, gcsafe.} =
    try:
        let path = req.url.path
        if routes.hasKey(path) and routes[path].hasKey($req.reqMethod):
            let (body, statusCode) = await routes[path][$req.reqMethod](req)
            await req.respond(statusCode, $body, newHttpHeaders({"Content-Type": "application/json"}))
        else:
            await req.respond(Http404, "Not Found")
    except Exception:
        await req.respond(Http500, getCurrentExceptionMsg())

route("POST", "/block/ip", block_ip)

cfg = getConfig()
let server = newAsyncHttpServer()
waitFor server.serve(Port(cfg.port), cb)
